'''
Author: zhanwei xu
Date: 2023-06-10 10:41:13
LastEditors: zhanwei xu
LastEditTime: 2023-06-12 10:21:12
Description: 

Copyright (c) 2023 by zhanwei xu, Tsinghua University, All Rights Reserved. 
'''
import pymongo as mg
import requests
import httpx
from typing import Any, Dict, List, Optional
from pydantic import BaseModel
from fastapi import APIRouter, Request, HTTPException, UploadFile, File
from fastapi.responses import JSONResponse
from datetime import datetime, timedelta
from core import setting, mgClient
from uuid import uuid1
from sys import getsizeof
db = mgClient['chat']
collection = db.get_collection('file')
current_timestamp = int(datetime.now().timestamp()*1000)
collection.insert_one({
    "ip": 0,
    "timestamp": {"$lt": current_timestamp}
})
router = APIRouter(
    prefix="/share",
    tags=["share"]
)

def check_ip_request_count(ip: str) -> bool:
    now = datetime.now()
    yesterday = now - timedelta(days=1)

    # Clean up IPs older than 24 hours
    db.requests.delete_many({"timestamp": {"$lt": yesterday}})

    # Get current IP requests count in the last 24 hours
    ip_requests = list(db.requests.find({"ip": ip, "timestamp": {"$gte": yesterday}}))
    ip_request_count = len(ip_requests)

    if ip_request_count < 30:
        # Insert a new request record if request count is less than 3
        db.requests.insert_one({"ip": ip, "timestamp": now})
        return True
    else:
        return False

@router.post("/upload_file/")
async def upload_file(request: Request, file: UploadFile = File(...)):

    client_ip = request.client.host
    if not check_ip_request_count(client_ip):
        raise HTTPException(status_code=429, detail="IP 请求次数超过限制，每24小时最多3次")
    file_size = getsizeof(file.file)
    if file_size > 1024*1024*10:
        raise HTTPException(status_code=400, detail="文件大小超过10MB")

    # `target_server_url` is a variable that stores the URL of the backend server to which the file
    # will be uploaded. It is used in the `upload_file` function to send the file to the backend
    # server using the `httpx` library.
    target_server_url = "http://www.superatmas.top:50001/upload_file/"
        
    async with httpx.AsyncClient() as client:
        # 读取文件内容以传输
        file_content = await file.read()
        
        # 发送文件到后端服务器
        response = await client.post(
            target_server_url,
            files={"file": (file.filename, file_content, file.content_type)},
        )
    # files = {"file": (file.filename, file_data, file.content_type)}
    # response = requests.post(target_server_url, files=files)

    if response.status_code != 200:
        raise HTTPException(status_code=response.status_code, detail=response.text)

    return JSONResponse(response.json())

@router.post("/message/")
async def message(request: Request) -> Optional[str]:
    json_data = await request.json()
    filename = json_data["filename"]
    question = json_data["question"]
    proxy_server = "http://www.superatmas.top:50001/message/"
    content = {"filename": filename, "question": question}
    response = requests.post(proxy_server, json=content)

    if response.status_code != 200:
        raise HTTPException(status_code=response.status_code, detail="Error forwarding to original server")

    json_response = response.json()
    return JSONResponse({"answer": json_response["answer"]})